import numpy as np

class ProbCalculator():
    def __init__(self,template):
        self.template=template
        self.guess_key_likelihood=np.zeros((256))

    def update(self,traces,midvalue):
        if len(midvalue.shape)==2:
            template = self.template.transpose()
            ln_likelihood  = np.sum((traces[:, :, None] - template) ** 2, axis=1)
            guess_key_likelihood=np.sum([ln_likelihood[i, midvalue[i]] for i in range(traces.shape[0])], axis=0)
            self.guess_key_likelihood=self.guess_key_likelihood+guess_key_likelihood
        else:
            tshape_2d = (self.template.shape[0] * 2, self.template.shape[2])
            template_2d = self.template.reshape(tshape_2d).transpose()
            ln_likelihood_2d = np.sum((traces[:, :, None] - template_2d) ** 2, axis=1)
            ln_likelihood_3d = ln_likelihood_2d.reshape((ln_likelihood_2d.shape[0], ln_likelihood_2d.shape[1] >> 1, 2))
            guess_key_likelihood=np.sum([np.sum([ln_likelihood_3d[i][j, midvalue[i].transpose()[j]] for j in range(midvalue.shape[2])],axis=0) for i in range(traces.shape[0])],axis=0)
            self.guess_key_likelihood = self.guess_key_likelihood + guess_key_likelihood

    def get_prob(self):
        return self.guess_key_likelihood


def profile(self):
    traces_shape = self.traces_profile.shape
    midvalue_shape = self.midvalue_profile.shape
    if len(midvalue_shape)==2:
        self.template=np.zeros((midvalue_shape[1],2,traces_shape[1]))
        for i in range(midvalue_shape[1]):
            index0=np.where(self.midvalue_profile[:,i]==0)[0]
            self.template[i,0]=np.mean(self.traces_profile[index0],axis=0)
            index1=np.where(self.midvalue_profile[:,i]==1)[0]
            self.template[i,1]=np.mean(self.traces_profile[index1],axis=0)
            yield 100*i/midvalue_shape[1]
    else:
        category=np.max(self.midvalue_profile)+1
        self.template = np.zeros((category, traces_shape[1]))
        for i in range(category):
            index=np.where(self.midvalue_profile==i)[0]
            if len(index)==0:
                continue
            self.template[i]=np.mean(self.traces_profile[index],axis=0)
            yield 100*i/category

def attack(self):
    trs_num = self.traces_attack.shape[0]
    attack_steps = self.param_dict["attack_steps"]
    index = [i for i in range(0, trs_num, attack_steps)]
    index.append(trs_num)
    prob_calculator=ProbCalculator(self.template)
    for i in range(len(index) - 1):
        prob_calculator.update(self.traces_attack[index[i]:index[i + 1]], self.midvalue_attack[index[i]:index[i + 1]])
        prob = prob_calculator.get_prob()
        yield prob, index[i + 1]

def get_guess_key(self, prob):
    return np.argmin(prob)

def get_rank(self,prob,real_key):
    prob_sort=np.argsort(np.argsort(prob))
    return prob_sort[real_key]
